package cn.janescott.framework.druid.servlet;

import com.alibaba.druid.stat.DruidStatService;
import com.alibaba.druid.support.http.ResourceServlet;
import com.alibaba.druid.support.http.StatViewServlet;
import com.alibaba.druid.support.logging.Log;
import com.alibaba.druid.support.logging.LogFactory;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import javax.management.MBeanServerConnection;
import javax.management.ObjectName;
import javax.management.remote.JMXConnector;
import javax.management.remote.JMXConnectorFactory;
import javax.management.remote.JMXServiceURL;
import javax.servlet.ServletException;
import java.io.IOException;
import java.io.Serializable;
import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 自定义监控servlet
 * 1、新增查询所有分布式数据源连接的url{@link /listJmx.json}
 * 2、新增切换当前监控数据源的url{@link /changeJmx.json?jmxName=xxx}
 *
 * @author scott
 * @date 2019/2/15
 */
public class MyStatViewServlet extends StatViewServlet {
    private final static Log LOG = LogFactory.getLog(StatViewServlet.class);

    /**
     * 修改分布式数据源请求路径
     */
    private static final String LIST_JMX = "/listJmx.json";
    private static final String CHANGE_JMX = "/changeJmx.json";
    private static final String JMX_NAME = "jmxName";
    /**
     * 分布式数据源配置
     */
    private static final String JMX_URLS = "jmxUrls";
    private static final String JMX_USER_NAMES = "jmxUserNames";
    private static final String JMX_PASS_WORDS = "jmxPassWords";
    private static final String JMX_DRUID_NAMES = "jmxDruidNames";

    /**
     * 初始化时检查数据源是否可用
     */
    private static final String JMX_INITIAL_CHECK = "jmxInitialCheck";
    /**
     * 检查不可用是否启动失败
     */
    private static final String JMX_FAILED_ON_CHECK = "jmxFailedOnCheck";

    /**
     * jmx连接名称
     */
    private AtomicInteger order = new AtomicInteger();
    /**
     * 修改jmx的锁
     */
    private final ReentrantLock lock = new ReentrantLock();
    /**
     * jmx的名字
     */
    private String jmxName;
    /**
     * jmx连接
     * 使用此连接时，需要获取{@see lock}锁
     * 因为此连接对象可能因为切换数据源导致失效
     */
    private MBeanServerConnection conn = null;

    /**
     * 分布式数据源列表
     */
    private final Map<String, JmxMetaData> jmxMetas = new LinkedHashMap<>();
    private final Map<String, String> jmxStatus = new LinkedHashMap<>();
    /**
     * jmx连接状态
     */
    private final String UP = "UP";
    private final String DOWN = "DOWN";


    /**
     * 尝试自定义初始化，如果失败，则调用父类初始化
     *
     * @throws ServletException
     */
    @Override
    public void init() throws ServletException {
        if (!initJmxConfig()) {
            super.init();
        }
    }

    /**
     * 适配原生逻辑 && 添加自己分布式数据源监控功能
     *
     * @param url
     * @return
     */
    @Override
    protected String process(String url) {
        if (CollectionUtils.isEmpty(jmxMetas)) {
            if (url.startsWith(LIST_JMX) || url.startsWith(CHANGE_JMX)) {
                return DruidStatService.returnJSONResult(DruidStatService.RESULT_CODE_SUCCESS, "not distributed datasource config!");
            }
            return super.process(url);
        }
        String resp = null;
        if (url.startsWith(LIST_JMX)) {
            Map<String, Object> info = new HashMap<>(3);
            info.put("list", jmxStatus);
            info.put("current", jmxName);
            info.put("changeUrl", "/changeJmx.json?jmxName=");
            return DruidStatService.returnJSONResult(DruidStatService.RESULT_CODE_SUCCESS, info);
        }
        if (url.startsWith(CHANGE_JMX)) {
            Map<String, String> parameters = DruidStatService.getParameters(url);
            if (parameters.containsKey(JMX_NAME)) {
                String name = parameters.get(JMX_NAME);
                if (!jmxStatus.containsKey(name)) {
                    return DruidStatService.returnJSONResult(DruidStatService.RESULT_CODE_ERROR, "datasource name cannot be found!");
                }
                boolean changeJmx = changeJmx(name);
                if (changeJmx) {
                    return DruidStatService.returnJSONResult(DruidStatService.RESULT_CODE_SUCCESS, "datasource switch succeed!");
                } else {
                    return DruidStatService.returnJSONResult(DruidStatService.RESULT_CODE_ERROR, "datasource switch failed!");
                }
            } else {
                return DruidStatService.returnJSONResult(DruidStatService.RESULT_CODE_ERROR, "datasource name cannot be absent!");
            }
        }
        /* 连接在初始化时创建失败 */
        if (conn == null) {
            try {
                /* 尝试重新连接 */
                boolean changeJmx = changeJmx(jmxName);
                if (!changeJmx) {
                    resp = DruidStatService.returnJSONResult(DruidStatService.RESULT_CODE_ERROR, "init jmx connection error, please switch datasource!");
                }
            } catch (Exception e) {
                LOG.error("init jmx connection error", e);
                resp = DruidStatService.returnJSONResult(DruidStatService.RESULT_CODE_ERROR,
                        "init jmx connection error" + e.getMessage());
            }
            /* 连接成功 */
            if (conn != null) {
                lock.lock();
                try {
                    resp = getJmxResultCopy(conn, url);
                } catch (Exception e) {
                    LOG.error("get jmx data error", e);
                    resp = DruidStatService.returnJSONResult(DruidStatService.RESULT_CODE_ERROR, "get data error:"
                            + e.getMessage());
                } finally {
                    lock.unlock();
                }
            }
        }
        /* 连接成功 */
        else {
            lock.lock();
            try {
                resp = getJmxResultCopy(conn, url);
            } catch (Exception e) {
                LOG.error("get jmx data error", e);
                resp = DruidStatService.returnJSONResult(DruidStatService.RESULT_CODE_ERROR,
                        "get data error" + e.getMessage());
            } finally {
                lock.unlock();
            }
        }
        return resp;
    }

    /**
     * 自定义初始化连接
     *
     * @throws IOException
     */
    private void customInitJmxConn(JmxMetaData jmxMetaData) throws IOException {
        if (jmxMetaData != null) {
            lock.lock();
            try {
                conn = initialJmxConn(jmxMetaData.getUrl(), jmxMetaData.getUsername(), jmxMetaData.getPassword());
            } finally {
                lock.unlock();
            }
        }
    }

    /**
     * 创建jmx连接
     *
     * @param jmxUrl
     * @param jmxUsername
     * @param jmxPassword
     * @return
     * @throws IOException
     */
    private MBeanServerConnection initialJmxConn(String jmxUrl, String jmxUsername, String jmxPassword) throws IOException {
        JMXServiceURL url = new JMXServiceURL(jmxUrl);
        Map<String, String[]> env = null;
        if (jmxUsername != null) {
            env = new HashMap<>(1);
            String[] credentials = new String[]{jmxUsername, jmxPassword};
            env.put(JMXConnector.CREDENTIALS, credentials);
        }
        JMXConnector jmxc = JMXConnectorFactory.connect(url, env);
        return jmxc.getMBeanServerConnection();
    }

    /**
     * 修改jmx配置
     *
     * @param key
     * @return
     */
    private boolean changeJmx(String key) {
        JmxMetaData jmxMetaData = jmxMetas.get(key);
        if (jmxMetaData != null) {
            lock.lock();
            try {
                try {
                    customInitJmxConn(jmxMetaData);
                    jmxStatus.put(key, UP);
                    jmxName = key;
                } catch (IOException ex) {
                    jmxStatus.put(key, DOWN);
                    LOG.error(ex.getMessage());
                    return false;
                }
                return true;
            } finally {
                lock.unlock();
            }
        }
        throw new RuntimeException("incorrect jmx name: " + key);
    }

    /**
     * 初始化Jmx配置
     * 尝试连接
     */
    private boolean initJmxConfig() {
        // 获取jmx的连接配置信息
        String jmxUrls = invokeReadInitParam(JMX_URLS);
        if (StringUtils.hasText(jmxUrls)) {
            String usernames = invokeReadInitParam(JMX_USER_NAMES);
            String passwords = invokeReadInitParam(JMX_PASS_WORDS);
            String druidNames = invokeReadInitParam(JMX_DRUID_NAMES);
            /* 初始化元数据 */
            String firstKey = initMetaData(jmxUrls, usernames, passwords, druidNames);
            /* 尝试连接 */
            changeJmx(firstKey);
            /* 初始化环境 */
            Class<ResourceServlet> resourceServletClass = ResourceServlet.class;
            try {
                Method initAuthEnv = resourceServletClass.getDeclaredMethod("initAuthEnv");
                initAuthEnv.setAccessible(true);
                initAuthEnv.invoke(this);
            } catch (Exception ex) {
                LOG.error(ex.getMessage());
                throw new RuntimeException(ex.getMessage());
            }
        }
        return false;
    }

    /**
     * 初始化连接参数
     * 返回第一个数据源
     *
     * @param jmxUrls
     * @param usernames
     * @param passwords
     * @param druidNames
     * @return
     */
    private String initMetaData(String jmxUrls, String usernames, String passwords, String druidNames) {
        /* 是否检查 */
        boolean initCheck = Boolean.valueOf(invokeReadInitParam(JMX_INITIAL_CHECK));
        /* 是否失败 */
        boolean failedOnCheck = Boolean.valueOf(invokeReadInitParam(JMX_FAILED_ON_CHECK));
        String[] urls = jmxUrls.split(",");
        String[] names = usernames != null ? usernames.split(",") : new String[]{};
        String[] pwds = passwords != null ? passwords.split(",") : new String[]{};
        String[] dnames = druidNames != null ? druidNames.split(",") : new String[]{};
        String firstKey = null;
        for (int i = 0; i < urls.length; i++) {
            JmxMetaData jmxMetaData = new JmxMetaData();
            jmxMetaData.setUrl(urls[i].trim());
            if (names.length > i) {
                jmxMetaData.setUsername(names[i].trim());
            } else {
                jmxMetaData.setUsername(null);
            }
            if (pwds.length > i) {
                jmxMetaData.setPassword(pwds[i].trim());
            } else {
                jmxMetaData.setPassword(null);
            }
            String key;
            if (dnames.length > i && StringUtils.hasText(dnames[i])) {
                key = dnames[i] + order.getAndIncrement();
            } else {
                key = "druid" + order.getAndIncrement();
            }
            /* 检查连接是否可用 */
            if (initCheck) {
                boolean unavailable = true;
                try {
                    MBeanServerConnection mBeanServerConnection =
                            initialJmxConn(jmxMetaData.getUrl(), jmxMetaData.getUsername(), jmxMetaData.getPassword());
                    if (mBeanServerConnection != null) {
                        unavailable = false;
                    }
                } catch (IOException ioe) {
                    LOG.error(ioe.getMessage());
                }
                /* 启动失败 */
                if (failedOnCheck && unavailable) {
                    throw new RuntimeException(String.format("jmx for url[%s] is unavailable", jmxMetaData.getUrl()));
                }
                jmxStatus.put(key, unavailable ? DOWN : UP);
            }
            jmxStatus.putIfAbsent(key, "UNKNOWN");
            jmxMetas.put(key, jmxMetaData);
            if (firstKey == null) {
                firstKey = key;
            }
        }
        return firstKey;
    }

    /**
     * 调用父类的读取初始化参数的方法
     */
    private String invokeReadInitParam(String name) {
        Class<StatViewServlet> statViewServletClass = StatViewServlet.class;
        try {
            Method readInitParam = statViewServletClass.getDeclaredMethod("readInitParam", String.class);
            readInitParam.setAccessible(true);
            return (String) readInitParam.invoke(this, name);
        } catch (Exception ex) {
            LOG.error(ex.getMessage());
            throw new RuntimeException(ex.getMessage());
        }
    }

    /**
     * 复制父类的获取jmx结果的方法
     */
    private String getJmxResultCopy(MBeanServerConnection connection, String url) throws Exception {
        ObjectName name = new ObjectName(DruidStatService.MBEAN_NAME);
        return (String) connection.invoke(name, "service", new String[]{url},
                new String[]{String.class.getName()});
    }

    /**
     * JMX元数据
     */
    private static class JmxMetaData implements Serializable {
        /**
         * 连接地址
         */
        private String url;
        /**
         * 用户名
         */
        private String username;
        /**
         * 密码
         */
        private String password;

        public String getUrl() {
            return url;
        }

        public void setUrl(String url) {
            this.url = url;
        }

        public String getUsername() {
            return username;
        }

        public void setUsername(String username) {
            this.username = username;
        }

        public String getPassword() {
            return password;
        }

        public void setPassword(String password) {
            this.password = password;
        }
    }
}
